import pandas as pd  # 导入Pandas库，用于数据处理
import torch  # 导入PyTorch库，用于深度学习模型
from transformers import BertTokenizer, BertForSequenceClassification, AdamW, get_linear_schedule_with_warmup  # 导入Transformers库中的BERT相关模块
from torch.utils.data import DataLoader, TensorDataset  # 导入数据加载器和数据集处理模块
from tqdm.auto import tqdm  # 导入tqdm库，用于进度条显示
import numpy as np  # 导入NumPy库，用于数值计算
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix  # 导入评估指标
import matplotlib.pyplot as plt  # 导入Matplotlib库，用于绘图
from torch import nn  # 导入PyTorch中的神经网络模块
from sklearn.model_selection import StratifiedKFold, train_test_split  # 导入分层K折交叉验证和数据集划分方法
import seaborn as sns  # 导入Seaborn库，用于高级绘图
from typing import Dict, List  # 导入类型注解模块
import json  # 导入json模块，用于处理json格式数据
import os  # 导入os模块，用于操作文件和路径
from typing import Tuple

# 配置
CONFIG = {
    "data_path": "all_data.csv",  # 数据文件路径
    "model_path": "../bert/models--bert-base-chinese/snapshots/c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f",  # 预训练模型路径
    "num_labels": 3,  # 分类的类别数
    "max_length": 128,  # 文本序列的最大长度
    "batch_size": 64,  # 每个批次的大小
    "learning_rate": 2e-7,  # 学习率
    "weight_decay": 0.01,  # 权重衰减，用于防止过拟合
    "num_epochs": 200,  # 训练的轮数
    "num_warmup_steps": 0,  # 学习率预热的步数
    "patience": 5,  # 提前停止的耐心值，若验证集表现未提高则停止
    "k_folds": 5,  # K折交叉验证的折数
    "seed": 42  # 随机种子，保证实验可复现
}

# 设置设备为GPU或CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 设置随机种子，确保结果可复现
torch.manual_seed(CONFIG["seed"])
np.random.seed(CONFIG["seed"])


# 加载数据
def load_data(path: str) -> pd.DataFrame:
    df = pd.read_csv(path)  # 读取CSV文件
    label_dict = {'positive': 0, 'negative': 1, 'neutral': 2}  # 标签映射
    df['label'] = df['label'].map(label_dict)  # 将标签映射为数字
    return df  # 返回数据框


# 数据预处理
def prepare_data(df: pd.DataFrame, tokenizer: BertTokenizer) -> Dict[str, torch.Tensor]:
    encodings = tokenizer(df['comment'].tolist(), truncation=True, padding=True,
                          max_length=CONFIG["max_length"])  # 对文本进行编码
    return {
        'input_ids': torch.tensor(encodings['input_ids']),  # 输入ID
        'attention_mask': torch.tensor(encodings['attention_mask']),  # 注意力掩码
        'labels': torch.tensor(df['label'].tolist())  # 标签
    }


# 创建数据加载器
def create_data_loader(dataset: TensorDataset, batch_size: int, shuffle: bool = True) -> DataLoader:
    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)  # 返回DataLoader对象


# 模型训练
# 模型训练
def train_model(model: nn.Module, tokenizer: BertTokenizer, train_loader: DataLoader, val_loader: DataLoader,
                optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler._LRScheduler, fold: int) -> Dict[str, List[float]]:
    best_val_f1 = 0  # 初始化最佳F1值
    no_improve = 0  # 无改进的轮次数
    history = {'train_loss': [], 'val_loss': [], 'train_f1': [], 'val_f1': [], 'train_accuracy': [], 'val_accuracy': []}  # 记录训练过程的损失和F1值

    for epoch in range(CONFIG["num_epochs"]):  # 训练多个epoch
        model.train()  # 模型进入训练模式
        train_loss, train_f1, train_acc, train_precision, train_recall = run_epoch(model, train_loader, optimizer, scheduler)  # 运行一个训练epoch

        model.eval()  # 模型进入评估模式
        val_loss, val_f1, val_acc, val_precision, val_recall = run_epoch(model, val_loader)  # 运行一个验证epoch

        history['train_loss'].append(train_loss)  # 记录训练损失
        history['val_loss'].append(val_loss)  # 记录验证损失
        history['train_f1'].append(train_f1)  # 记录训练F1分数
        history['val_f1'].append(val_f1)  # 记录验证F1分数
        history['train_accuracy'].append(train_acc)  # 记录训练准确率
        history['val_accuracy'].append(val_acc)  # 记录验证准确率

        # 打印当前epoch的训练和验证指标
        print(
            f"Epoch {epoch + 1}/{CONFIG['num_epochs']}, "
            f"Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, "
            f"Train Accuracy: {train_acc:.4f}, Val Accuracy: {val_acc:.4f}, "
            f"Train Precision: {train_precision:.4f}, Val Precision: {val_precision:.4f}, "
            f"Train Recall: {train_recall:.4f}, Val Recall: {val_recall:.4f}, "
            f"Train F1-score: {train_f1:.4f}, Val F1-score: {val_f1:.4f}"
        )

        if val_f1 > best_val_f1:  # 如果当前验证集F1分数是最好的
            best_val_f1 = val_f1  # 更新最佳F1分数
            # 保存模型和分词器
            model.save_pretrained(f"./sentiment_best_model_accuracy/fold_{fold}")  # 保存模型
            tokenizer.save_pretrained(f"./sentiment_best_model_accuracy/fold_{fold}")  # 保存分词器
            no_improve = 0  # 重置无改进轮数
        else:
            no_improve += 1  # 无改进轮数加1

        if no_improve >= CONFIG["patience"]:  # 若无改进的轮数达到提前停止的耐心值
            print(f"Early stopping triggered after epoch {epoch + 1}")  # 打印提前停止信息
            break

    return history  # 返回训练历史


# 运行一个epoch，增加指标计算
def run_epoch(model: nn.Module, data_loader: DataLoader, optimizer: torch.optim.Optimizer = None,
              scheduler: torch.optim.lr_scheduler._LRScheduler = None) -> Tuple[float, float, float, float, float]:
    total_loss = 0  # 总损失初始化为0
    all_preds = []  # 用于存储所有的预测值
    all_labels = []  # 用于存储所有的真实标签

    for batch in tqdm(data_loader, desc="Processing batches"):  # 遍历数据批次，显示进度条
        input_ids = batch[0].to(device)  # 将输入ID移到设备（GPU/CPU）
        attention_mask = batch[1].to(device)  # 将注意力掩码移到设备
        labels = batch[2].to(device)  # 将标签移到设备

        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)  # 模型前向传播，计算输出
        loss = outputs.loss  # 计算损失

        if optimizer:  # 如果是训练过程
            optimizer.zero_grad()  # 梯度清零
            loss.backward()  # 反向传播计算梯度
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 梯度裁剪
            optimizer.step()  # 更新模型参数
            if scheduler:
                scheduler.step()  # 更新学习率调度器

        total_loss += loss.item()  # 累加损失
        preds = torch.argmax(outputs.logits, dim=1)  # 计算预测值
        all_preds.extend(preds.cpu().numpy())  # 将预测值移到CPU并添加到列表
        all_labels.extend(labels.cpu().numpy())  # 将标签移到CPU并添加到列表

    avg_loss = total_loss / len(data_loader)  # 计算平均损失

    # 计算各项性能指标
    accuracy = accuracy_score(all_labels, all_preds)  # 计算准确率
    precision = precision_score(all_labels, all_preds, average='macro')  # 计算平均精确率
    recall = recall_score(all_labels, all_preds, average='macro')  # 计算宏、平均召回率
    f1 = f1_score(all_labels, all_preds, average='macro')  # 计算宏、平均F1分数

    return avg_loss, f1, accuracy, precision, recall  # 返回平均损失、F1分数、准确率、精确率、召回率




# 主函数
def main():
    # 加载并预处理数据
    df = load_data(CONFIG["data_path"])  # 加载数据
    df_train, df_test = train_test_split(df, test_size=0.2, stratify=df['label'],
                                         random_state=CONFIG["seed"])  # 划分训练集和测试集
    df_test.to_csv("test_data.csv", index=False)  # 将测试集保存为CSV文件
    tokenizer = BertTokenizer.from_pretrained(CONFIG["model_path"])  # 加载BERT的分词器

    # 准备测试数据
    test_data = prepare_data(df_test, tokenizer)  # 预处理测试数据
    test_dataset = TensorDataset(test_data['input_ids'], test_data['attention_mask'], test_data['labels'])  # 创建测试数据集
    test_loader = create_data_loader(test_dataset, CONFIG["batch_size"], shuffle=False)  # 创建测试数据加载器

    # 初始化分层K折交叉验证
    skf = StratifiedKFold(n_splits=CONFIG["k_folds"], shuffle=True, random_state=CONFIG["seed"])  # 初始化StratifiedKFold

    for fold, (train_idx, val_idx) in enumerate(skf.split(df_train, df_train['label']), 1):  # 遍历每一折
        print(f"Fold {fold}/{CONFIG['k_folds']}")  # 打印当前折数

        # 准备当前折的训练和验证数据
        df_train_fold = df_train.iloc[train_idx]  # 获取当前折的训练数据
        df_val_fold = df_train.iloc[val_idx]  # 获取当前折的验证数据

        train_data = prepare_data(df_train_fold, tokenizer)  # 预处理训练数据
        val_data = prepare_data(df_val_fold, tokenizer)  # 预处理验证数据

        train_dataset = TensorDataset(train_data['input_ids'], train_data['attention_mask'],
                                      train_data['labels'])  # 创建训练数据集
        val_dataset = TensorDataset(val_data['input_ids'], val_data['attention_mask'], val_data['labels'])  # 创建验证数据集

        train_loader = create_data_loader(train_dataset, CONFIG["batch_size"])  # 创建训练数据加载器
        val_loader = create_data_loader(val_dataset, CONFIG["batch_size"], shuffle=False)  # 创建验证数据加载器

        # 初始化模型和优化器
        model = BertForSequenceClassification.from_pretrained(CONFIG["model_path"],
                                                              num_labels=CONFIG["num_labels"])  # 加载预训练模型
        model.to(device)  # 将模型移到设备上

        optimizer = AdamW(model.parameters(), lr=CONFIG["learning_rate"],
                          weight_decay=CONFIG["weight_decay"])  # 初始化AdamW优化器
        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=CONFIG["num_warmup_steps"],
                                                    num_training_steps=len(train_loader) * CONFIG[
                                                        "num_epochs"])  # 初始化学习率调度器

        # 训练模型
        history = train_model(model, tokenizer, train_loader, val_loader, optimizer, scheduler, fold)  # 训练并返回训练历史

        # 绘制训练历史图
        plot_training_history(history, fold)

    # 集成预测测试集
    ensemble_prediction(test_loader)


# 绘制训练历史图
def plot_training_history(history: Dict[str, List[float]], fold: int):
    plt.figure(figsize=(12, 8))  # 设置图像大小
    plt.subplot(2, 1, 1)  # 第一个子图
    plt.plot(history['train_loss'], label='Train Loss')  # 绘制训练损失曲线
    plt.plot(history['val_loss'], label='Validation Loss')  # 绘制验证损失曲线
    plt.title(f'Training and Validation Loss (Fold {fold})')  # 设置图像标题
    plt.legend()  # 显示图例

    plt.subplot(2, 1, 2)  # 第二个子图
    plt.plot(history['train_f1'], label='Train F1')  # 绘制训练F1曲线
    plt.plot(history['val_f1'], label='Validation F1')  # 绘制验证F1曲线
    plt.title(f'Training and Validation F1 Score (Fold {fold})')  # 设置图像标题
    plt.legend()  # 显示图例

    plt.tight_layout()  # 调整布局
    plt.savefig(f'training_history_fold_{fold}.png')  # 保存图像
    plt.close()  # 关闭图像


# 测试集集成预测
def ensemble_prediction(test_loader: DataLoader):
    models = []  # 用于存储每一折的模型
    for i in range(1, CONFIG["k_folds"] + 1):  # 遍历每一折
        # 直接加载保存的模型，而不是仅加载权重
        model = BertForSequenceClassification.from_pretrained(f"./sentiment_best_model_accuracy/fold_{i}",
                                                              num_labels=CONFIG["num_labels"])  # 加载保存的模型
        model.to(device)  # 将模型移到设备
        model.eval()  # 设置模型为评估模式
        models.append(model)  # 将模型添加到模型列表中

    all_preds = []  # 存储所有的预测结果
    all_labels = []  # 存储所有的真实标签

    with torch.no_grad():  # 不计算梯度
        for batch in tqdm(test_loader, desc="Testing"):  # 遍历测试集批次
            input_ids = batch[0].to(device)  # 从元组中提取input_ids并移到设备
            attention_mask = batch[1].to(device)  # 从元组中提取attention_mask并移到设备
            labels = batch[2].to(device)  # 从元组中提取labels并移到设备
            all_labels.extend(labels.cpu().numpy())  # 将标签移动到CPU并保存

            fold_preds = []  # 存储当前批次每个模型的预测
            for model in models:  # 遍历每一个模型
                outputs = model(input_ids, attention_mask=attention_mask)  # 模型前向传播
                preds = torch.argmax(outputs.logits, dim=1)  # 获取模型预测的类别
                fold_preds.append(preds.cpu().numpy())  # 保存当前模型的预测结果

            fold_preds = np.array(fold_preds)  # 转换为NumPy数组
            ensemble_preds = np.mean(fold_preds, axis=0).round().astype(int)  # 计算每个样本的平均预测，并四舍五入

            all_preds.extend(ensemble_preds)  # 将集成预测结果保存

    # 计算性能指标
    acc = accuracy_score(all_labels, all_preds)  # 计算准确率
    precision = precision_score(all_labels, all_preds, average='macro')  # 计算宏平均精确率
    recall = recall_score(all_labels, all_preds, average='macro')  # 计算宏平均召回率
    f1 = f1_score(all_labels, all_preds, average='macro')  # 计算宏平均F1分数

    # 打印结果
    print(f"Test Accuracy: {acc:.4f}")  # 打印测试集准确率
    print(f"Test Precision: {precision:.4f}")  # 打印测试集精确率
    print(f"Test Recall: {recall:.4f}")  # 打印测试集召回率
    print(f"Test F1 Score: {f1:.4f}")  # 打印测试集F1分数

    # 绘制混淆矩阵
    cm = confusion_matrix(all_labels, all_preds)  # 计算混淆矩阵
    plt.figure(figsize=(8, 6))  # 设置图像大小
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')  # 绘制热力图
    plt.title('Confusion Matrix')  # 设置标题
    plt.xlabel('Predicted')  # 设置x轴标签
    plt.ylabel('True')  # 设置y轴标签
    plt.show()  # 显示图像



if __name__ == "__main__":
    main()  # 运行主函数